import backtrader as bt
from .DbOprate import DbOprate

class TDXSMA(bt.Indicator):
    """
    通达信SMA指标实现: SMA(X,N,M) = (X*M + Y'*(N-M)) / N
    其中Y'是前一日的SMA值
    """
    lines = ('sma',)
    params = (
        ('period', 2),
        ('weight', 1),
    )

    def __init__(self):
        super(TDXSMA, self).__init__()
        self.addminperiod(self.p.period)

    def next(self):
        if len(self) == 1:
            # 第一天，使用当前值作为初始值
            self.lines.sma[0] = self.data[0]
        else:
            # SMA = (current_value * weight + previous_sma * (period - weight)) / period
            current_val = self.data[0]
            previous_sma = self.lines.sma[-1]
            
            # 检查previous_sma是否为nan
            if previous_sma != previous_sma:  # nan检查
                self.lines.sma[0] = current_val
            else:
                self.lines.sma[0] = (current_val * self.p.weight + 
                                   previous_sma * (self.p.period - self.p.weight)) / self.p.period

class TDXZTStrategy(bt.Strategy):
    """
    涨停 := C >= REF(C, 1) * 1.098 AND C = H;
    量能非最高 := V < HHV(V, 20);
    首次涨停 := COUNT(涨停, 20) = 1 AND 涨停;
    次日放量跌 := REF(首次涨停, 1) AND C < O AND V > REF(V, 1) AND C<REF(C,1);
    买点 : REF(次日放量跌, 1) AND C > O AND C > REF(C, 1) AND O=L AND O<REF(C,1);
    """
    params = (
        ('hold_days', 5),        # 持有天数
        ('profit_target', 0.08), # 盈利目标 8%
        ('stop_loss', 0.05),     # 止损 5%
        ('position_ratio', 0.9), # 使用资金比例 90%
        ('filter_period', 10),   # 信号过滤周期
    )

    def log(self, txt, dt=None):
        dt = dt or self.datas[0].datetime.date(0)
        print(f'{dt.isoformat()}, {txt}')

    def __init__(self):
        self.data_close = self.datas[0].close

        # X_7 := MA(CLOSE,1)  实际上就是收盘价
        self.x_7 = self.data_close
        
        # X_8 := SMA(SMA(SMA(X_7,2,1),2,1),2,1)
        self.x_8_step1 = TDXSMA(self.x_7, period=2, weight=1)          
        self.x_8_step2 = TDXSMA(self.x_8_step1, period=2, weight=1)    
        self.x_8 = TDXSMA(self.x_8_step2, period=2, weight=1)          
        
        # X_9 := MA(X_8,3) 
        self.x_9 = bt.ind.SMA(self.x_8, period=3)
        
        # HL := X_8 >= X_9
        self.hl = self.x_8 >= self.x_9
        
        # LL := X_8 <= X_9  
        self.ll = self.x_8 <= self.x_9
        
        # ZZD := HL AND REF(LL, 1)
        self.zzd = bt.And(self.hl, self.ll(-1))
        
        # LXYZ := COUNT(LL, 10) = 9
        self.ll_count = bt.ind.SumN(self.ll, period=10)
        self.lxyz = (self.ll_count == 9)
        
        # B := ZZD AND LXYZ (买入信号)
        self.buy_signal = bt.And(self.zzd, self.lxyz)
        
        # 其他变量
        self.filter_counter = 0
        self.last_signal_bar = -1
        self.buy_price = None
        self.buy_date = None
        self.hold_bars = 0

    def next(self):
        current_date = self.datas[0].datetime.date(0)
        current_price = self.data_close[0]

        # 获取资金状态
        cash = max(0, self.broker.getcash())
        position_size = self.position.size
        position_value = position_size * current_price if position_size > 0 else 0
        
        # 更新过滤计数器
        if self.filter_counter > 0:
            self.filter_counter -= 1

        # 持仓管理逻辑
        if position_size > 0 and self.buy_price is not None:
            self.hold_bars += 1
            
            # 检查卖出条件
            should_sell = False
            sell_reason = ""
            
            # 止损止盈逻辑
            if self.hold_bars >= self.p.hold_days:
                should_sell = True
                sell_reason = f"持有满{self.p.hold_days}天"
            elif (current_price - self.buy_price) / self.buy_price >= self.p.profit_target:
                should_sell = True
                sell_reason = f"盈利达到{self.p.profit_target*100}%"
            elif (self.buy_price - current_price) / self.buy_price >= self.p.stop_loss:
                should_sell = True
                sell_reason = f"亏损达到{self.p.stop_loss*100}%"
            
            # 执行卖出
            if should_sell:
                size = position_size
                
                self.log(f'*** 卖出信号 *** 原因: {sell_reason}, 价格: {current_price:.2f}, 数量: {size}')
                
                # 执行卖出
                self.sell(size=size)
                
                # 重置变量
                self.buy_price = None
                self.buy_date = None
                self.hold_bars = 0
        else:
            # 检查买入信号 - 确保所有指标都有有效值
            if (len(self.x_8) > 0 and len(self.x_9) > 0 and 
                len(self.buy_signal) > 0 and
                self.x_8[0] == self.x_8[0] and self.x_9[0] == self.x_9[0] and  # 检查nan
                self.filter_counter <= 0 and self.buy_signal[0] and cash > 1000):
                
                self.log(f'*** 买入信号触发 ***')
                self.log(f'*** 可用资金: {cash:.2f}, 当前价格: {current_price:.2f}')
                self.log(f'X_8: {self.x_8[0]:.6f}, X_9: {self.x_9[0]:.6f}')
                
                # 计算买入数量
                use_cash = cash * self.p.position_ratio
                max_size = int((use_cash / current_price) / 100) * 100
                
                if max_size >= 100:
                    size = max_size
                    
                    # 存入数据库 - 买入信号
                    try:
                        result = {
                            'strategy': 'TDXZTStrategy',
                            'stock_code': getattr(self.datas[0], '_name', 'UNKNOWN'),
                            'op_act': 'buy_signal',
                            'op_price': float(current_price),
                            'op_num': int(size),
                            'op_comm': 0.0,
                            'op_dt': current_date.isoformat(),
                            'reason': 'TDX_BUY_SIGNAL',
                            'can_use_money': float(max(0, cash)),
                            'init_cash': float(getattr(self.broker, 'startingcash', cash)),
                            'stock_num': int(max(0, position_size)),
                            'stock_amount': float(max(0, position_value)),
                            'profitable': 0.0,
                            'x_8': float(self.x_8[0]) if self.x_8[0] == self.x_8[0] else 0.0,
                            'x_9': float(self.x_9[0]) if self.x_9[0] == self.x_9[0] else 0.0
                        }
                        
                        # 调用数据库操作保存结果
                        DbOprate.save_result_to_mysql(result)
                        self.log(f'*** 买入信号已存入数据库 *** 价格: {current_price:.2f}, 数量: {size}')
                        
                        # 启动过滤期
                        self.filter_counter = self.p.filter_period
                        
                    except Exception as e:
                        self.log(f'*** 数据库存储失败 *** 错误: {str(e)}')

    def notify_order(self, order):
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f'*** 买入订单完成 *** 价格: {order.executed.price:.2f}, 数量: {order.executed.size}')
                self.buy_price = order.executed.price
                self.buy_date = self.datas[0].datetime.date(0)
                self.hold_bars = 0
            elif order.issell():
                self.log(f'*** 卖出订单完成 *** 价格: {order.executed.price:.2f}, 数量: {order.executed.size}')
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log(f'*** 订单失败 *** 状态: {order.Status[order.status]}')

    def notify_trade(self, trade):
        if trade.isclosed:
            profit = trade.pnl
            self.log(f'*** 交易完成 *** 盈亏: {profit:.2f}')
